/**
 * @file FEMSpace.h
 * @author Lishijie (lsj1018845759@outlook.com)
 * @brief 
 * @version 0.1
 * @date 2020-11-11
 * 
 * @copyright Copyright (c) 2020
 * 
 */
#ifndef _FEMSPACE_H_
#define _FEMSPACE_H_
#include "Mesh.h"
#include "Element.h"
#include <iostream>
#include <map>
#include "Equation.h"
#include <Eigen/Sparse>
#include "BoundaryFunction.h"
#include "Eigen/IterativeLinearSolvers"

#define TEMPLATE template<Dimension DIM>
typedef Eigen::SparseMatrix<Real> SpMat;
typedef Eigen::Triplet<Real> Tri;
typedef Eigen::VectorXd RHS;
typedef Eigen::VectorXd Solution;
/**
 * @brief 有限元空间基类，各种当方程的求解可以继承该类。
 *        提供的功能有：组装刚度矩阵，组装右端项，边界条件处理，计算数值解并返回；
 * 
 */
TEMPLATE
class FEMSpace
{
public:
    /**
     * @brief 组装刚度矩阵；
     */
    virtual void AssembleStiffMatrix() = 0;
    /**
     * @brief 组装右端项；
     * 
     */
    virtual void AssembleRightHandsTerm() = 0;
    /**
     * @brief 处理边界条件
     * 
     */
    virtual void DealWithBoubdaryCondition() = 0;
    /**
     * @brief 求解
     * 
     */
    virtual void Solve() = 0;
    /**
     * @brief 返回解
     * 
     * @return Solution 
     */
    virtual Solution solution() = 0;
    /**
     * @brief 局部刚度矩阵在i,j处的值
     * 
     * @param i 下标１
     * @param j 下标２
     * @return Real 值
     */
    virtual Real LocalStiffMatrix(int i, int j) = 0;
};
#undef TEMPLATE

class Possion_2D:public FEMSpace<2>
{


public:
    Possion_2D();
    Possion_2D(Mesh<2>*,Element<2>*,Equation<2>*,BoundaryCondition<2>);
    Possion_2D(const Possion_2D& );
    ~Possion_2D(){/*delete _mesh;delete _element;delete _eqation;*/};
    virtual void AssembleStiffMatrix();
    virtual void AssembleRightHandsTerm();
    virtual void DealWithBoubdaryCondition();
    virtual void Solve();
    virtual Solution solution();
    virtual Real LocalStiffMatrix(int i, int j);
    SpMat A() const;
    RHS Rhs() const;
protected:
    SpMat _A;
    RHS _Rhs;
    Solution _Solution;
    Mesh<2>* _mesh;
    Element<2>* _element;
    Equation<2> *_eqation;
    BoundaryCondition<2> _boundarycondition;
    
};
Possion_2D::Possion_2D()
{   
    _A = SpMat(int(_mesh->n_dofs()), int(_mesh->n_dofs()));
    _Rhs = RHS(int(_mesh->n_dofs()));
    _boundarycondition.apply(_mesh);
}
Possion_2D::Possion_2D(Mesh<2>* mesh,Element<2>* element,Equation<2>*equation,BoundaryCondition<2> BoundaryCondition)
{
    _mesh = mesh;
    _element = element;
    _eqation = equation;
    _boundarycondition = (BoundaryCondition);
    _A = SpMat(int(_mesh->n_dofs()), int(_mesh->n_dofs()));
    _Rhs = RHS(int(_mesh->n_dofs()));
    _boundarycondition.apply(_mesh);
}

Possion_2D::Possion_2D(const Possion_2D& p)
{
    _mesh = p._mesh;
    _element = p._element;
    _eqation = p._eqation;
    _boundarycondition = p._boundarycondition;
    _A = p._A;
    _Rhs = p._Rhs;
    _boundarycondition.apply(_mesh);
}

void Possion_2D::AssembleStiffMatrix()
{
    int n_Dofs = _element->n_Dofs();
    std::vector<Tri> TriList((_mesh->n_element()) * n_Dofs * n_Dofs);
    std::vector<Tri>::iterator it = TriList.begin(); 
    /// traverse each element
    for (int k = 0; k < _mesh -> n_element(); k++)
    {
        std::vector<int> Idx = _mesh->NodeofEle(k);
        std::vector<Dofs<2> > temp(n_Dofs);
        for(int i = 0;i < n_Dofs; i++)
        {
	    //temp[i].GlobalIndex() = Idx[i];
	    temp[i] = _mesh->DofsofIndex(Idx[i]);
        }
        _element->SetDofsList(temp);
        for (int i = 1; i <= n_Dofs; i++)
	    for (int j = 1; j <= n_Dofs; j++)
	    {
		*it = Tri(_element->NdIdx(i), _element->NdIdx(j), LocalStiffMatrix(i, j));
		it++;
	    }
    }
     _A.setFromTriplets(TriList.begin(), TriList.end());
     _A.makeCompressed();
     //std::cout << _A << std::endl;
}

void Possion_2D::AssembleRightHandsTerm()
{
    int n_Dofs = _element->n_Dofs();
    int n_GaussPnt = _element->n_GaussPnt();
    auto _func  = _eqation->RightHandsTermFunction();
    for (int k = 0; k < _mesh -> n_element(); k++)
    {
        std::vector<int> Idx = _mesh->NodeofEle(k);
        std::vector<Dofs<2>> temp(n_Dofs);
        for(int i = 0; i < n_Dofs; i++)
        {
	    //temp[i].GlobalIndex() = Idx[i];
	    temp[i] = _mesh ->DofsofIndex(Idx[i]);
        }
        _element->SetDofsList(temp);
        for(int i = 1; i <= n_Dofs; i++)
        {
            Real a = 0.0;
            for (int j = 0; j < n_GaussPnt; j++)
            {
		        double xi = _element->GaussionPoint(j)[0];//_GaussPnt[j][0];
		        double eta = _element->GaussionPoint(j)[1];
		        double xi_ = _element->Global_x(xi, eta);
		        double eta_ = _element->Global_y(xi, eta);
                Dofs<2> temp({xi_,eta_});
		        a = a + _element->det_Jacobi(xi, eta) * _element->GaussionWeight(j) * _element->phi(xi, eta, i)* _func(*temp);
            }
            _Rhs[_element->NdIdx(i)]+= a;
        }
    }
    //std::cout << _Rhs << std::endl;
    if(_boundarycondition.is_all_dirichlet())
        return;
    else
    {
        std::vector<BoundaryFunction<2>* > BFL = _boundarycondition.BFlist();
        std::vector<std::vector<int> > BndEdge = _mesh->_BndEdge;
        int num_of_bnd_edge = BndEdge.size();
        int n_Dofs = _element ->n_Dofs();
        int n_GaussPnt = _element ->n_GaussPnt();
        for(int k = 0;k < num_of_bnd_edge;k++)
        {
            std::function<double(double*)> _func;
            int bndtype = BndEdge[k][0];
            if(bndtype == -2)
            {
            int _element_index = BndEdge[k][1];
            std::vector<int> Idx = _mesh->NodeofEle(BndEdge[k][1]);
            std::vector<Dofs<2>> temp(n_Dofs);
            for(int i = 0; i < n_Dofs; i++)
	            temp[i] = _mesh ->DofsofIndex(Idx[i]);
            _element->SetDofsList(temp);
            int num1 = BndEdge[k][2];
            int num2 = BndEdge[k][3];
            Dofs<2> dof1 = _mesh->DofsofIndex(num1);
            Dofs<2> dof2 = _mesh->DofsofIndex(num2);
            bool is_horizontal = false;
            bool is_vertical = false;
            if(dof1[1] == dof2[1])
                is_horizontal = true;
            else if(dof1[0] == dof2[0])
                is_vertical = true;
            for(int i = 0;i < BFL.size();i++)
            {
                BoundaryFunction<2> * t = BFL[i];
                for(int j = 0;j < t->Boundaryvector().size();j++)
                {
                    Boundary<2> B = t->Boundaryvector()[j];
                    if(B.is_on_boundary(dof1) && B.is_on_boundary(dof2))
                    {
                        _func = t->g();
                        break;
                    }
                }
            }
            if(is_horizontal)
            {
                int c = dof1[1];
                bool is_bottom = (c == _mesh->RDomain()->LeftBottomPoint()[1]);
                for(int i = 1; i <= _element->n_Dofs(); i++)
                {
                Real a = 0.0;
                for (int j = 0; j < n_GaussPnt; j++)
                {
		            double xi = _element->GaussionPoint(j)[0];
		            double eta = _element->GaussionPoint(j)[1];
		            double xi_ = _element->Global_x(xi, eta);
                    Dofs<2> temp({xi_,double(c)});
                    int c_ = is_bottom ? -1 : 1;
		            a = a + sqrt(_element->det_Jacobi(xi, eta)) * _element->GaussionWeight(j) * _element->phi(xi, double(c_), i)* _func(*temp);
                }
                _Rhs[_element->NdIdx(i)]+= a / 2.0;
                }
            }
            else if(is_vertical)
            {
                int c = dof1[0];
                bool is_left = (c == _mesh->RDomain()->LeftBottomPoint()[0]);
                for(int i = 1; i <= _element->n_Dofs(); i++)
                {
                Real a = 0.0;
                for (int j = 0; j < n_GaussPnt; j++)
                {
		            double xi = _element->GaussionPoint(j)[0];
		            double eta = _element->GaussionPoint(j)[1];
		            double eta_ = _element->Global_y(xi, eta);
                    Dofs<2> temp({double(c),eta_});
                    int c_ = is_left ? -1 : 1;
		            a = a + sqrt(_element->det_Jacobi(xi, eta)) * _element->GaussionWeight(j) * _element->phi(double(c_), eta, i)* _func(*temp);
                }
                _Rhs[_element->NdIdx(i)]+= a / 2.0;
                }
            }
            }
        }
    }
    //std::cout << "neumann after" << std::endl;
    //std::cout << _Rhs << std::endl;

}

void Possion_2D::DealWithBoubdaryCondition()
{
    std::vector<BoundaryFunction<2>* > BFL = _boundarycondition.BFlist();
    if(_boundarycondition.have_dirichlet())
    for(auto k:_mesh->Boundary())
    {
        bool found_dirichlet_bnd = false;
        Dofs<2> bnd_point = _mesh ->DofsofIndex(k);
        Real bnd_value ;
        Boundaryinfo bctype;
        Boundaryinfo bctype_for_n;
        for(int i = 0; i < BFL.size();i++)
        {
            BoundaryFunction<2> * t = BFL[i];
            if(!found_dirichlet_bnd)
            {
            bctype =BFL[i]->boundaryType();           
            for(int j = 0;j < BFL[i]->Boundaryvector().size();j++)
            {
                Boundary<2> B =  BFL[i]->Boundaryvector()[j];
                if(B.is_on_boundary(bnd_point))
                {
                    if(bctype != 1)
                    {
                        bctype_for_n = bctype;
                        break;
                    }
                    else
                    {   
                        bctype_for_n = bctype;
                        bnd_value = BFL[i]->g()(*bnd_point);
                        found_dirichlet_bnd = true;
                    }
                }
            }
            }
        }
        if(bctype_for_n == 1)
        {
            _Rhs[k] = bnd_value * _A.coeffRef(k,k);
            for(Eigen::SparseMatrix<Real>::InnerIterator it(_A,k);it;++it)
            {
                int row = it.row();
                if(row == k)
                    continue;
                _Rhs[row] -= _A.coeffRef(k,row)* bnd_value;
                _A.coeffRef(k,row) = 0.0;
                _A.coeffRef(row,k) = 0.0;
            }
        }
    }
    //std::cout << _A << std::endl;
    //std::cout << _Rhs << std::endl;
}
void Possion_2D::Solve()
{
    Eigen::ConjugateGradient<Eigen::SparseMatrix<double> > Solver_sparse;
    Eigen::VectorXd x1_sparse;
    Solver_sparse.setTolerance(1e-12);
	Solver_sparse.compute(_A);
    _Solution = Solver_sparse.solve(_Rhs);
    //std::cout << "solution is : " << std::endl;
    //std::cout << _Solution << std::endl;
}

SpMat Possion_2D::A() const
{
    return _A;
}
RHS Possion_2D::Rhs() const
{
    return _Rhs;
}


Solution Possion_2D::solution()
{
    std::cout<< "The solution of this possion equation is : "<<std::endl;
     return _Solution;
}

Real Possion_2D::LocalStiffMatrix(int i, int j)
{
    Real det = _element -> det_Jacobi(0,0);
    Real a = 0;
    for (int k = 0; k < _element->n_GaussPnt(); k++)
    {
	Real xi = _element->GaussionPoint(k)[0];
	Real eta = _element->GaussionPoint(k)[1];
	//Real xx = _element->phi_x(xi, eta, i) * _element->phi_x(xi, eta, j);
	//Real yy = _element->phi_y(xi, eta, i) * _element->phi_y(xi, eta, j);
    a = a + _element->det_Jacobi(xi, eta) * _element->GaussionWeight(k) * InnerProuduct(_element->gradient(xi,eta,i),_element->gradient(xi,eta,j));
	//a = a + _element->det_Jacobi(xi, eta) * _element->GaussionWeight(k) * (xx + yy);
    }
    return static_cast<double>(a); 
}
#else
#endif